#!/usr/bin/python

'''
    OpenCV Image Alignment    
'''

import cv2
import numpy as np


def align_two_channel(im1_gray,im2_gray):
    # Convert images to grayscale
    #im1_gray = cv2.cvtColor(im1,cv2.COLOR_BGR2GRAY)
    #im2_gray = cv2.cvtColor(im2,cv2.COLOR_BGR2GRAY)
    
    # Find size of image1
    sz = im1.shape
 
    # Define the motion model
    warp_mode = cv2.MOTION_TRANSLATION
    

    # Define 2x3 or 3x3 matrices and initialize the matrix to identity
    if warp_mode == cv2.MOTION_HOMOGRAPHY :
        warp_matrix = np.eye(3, 3, dtype=np.float32)
    else :
        warp_matrix = np.eye(2, 3, dtype=np.float32)

    # Specify the number of iterations.
    number_of_iterations = 5000;
    
    # Specify the threshold of the increment
    # in the correlation coefficient between two iterations
    termination_eps = 1e-10;
    
    # Define termination criteria
    criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, number_of_iterations,  termination_eps)

    # Run the ECC algorithm. The results are stored in warp_matrix.
    #retval, warpMatrix	=	cv.findTransformECC(	templateImage, inputImage, warpMatrix[, motionType[, criteria[, inputMask]]]	)

    (cc, warp_matrix) = cv2.findTransformECC(im1_gray,im2_gray,warp_matrix,warp_mode, criteria,None, 5)
    


    if warp_mode == cv2.MOTION_HOMOGRAPHY :
        # Use warpPerspective for Homography
        im2_aligned = cv2.warpPerspective (im2, warp_matrix, (sz[1],sz[0]), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)
    else :
        # Use warpAffine for Translation, Euclidean and Affine
        im2_aligned = cv2.warpAffine(im2, warp_matrix, (sz[1],sz[0]), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP);

    print(cc)
    print(warp_matrix)
    return cc,warp_matrix,im2_aligned



# Read the images to be aligned
im1 =  cv2.imread("./img/5ch/IMG_0000_1.tif");
im2 =  cv2.imread("./img/5ch/IMG_0000_2.tif");
im3 =  cv2.imread("./img/5ch/IMG_0000_3.tif");

b=cv2.split(im1)[0];
g=cv2.split(im2)[0];
r=cv2.split(im3)[0];

rows,cols=g.shape

res_d = cv2.merge([b,g,r])

cc,wa,img=align_two_channel(b,g)
B1=cv2.warpAffine(b,wa,(cols,rows))
cc,wa,img=align_two_channel(r,g)
R1=cv2.warpAffine(r,wa,(cols,rows))

res_a = cv2.merge([B1,g,R1])

cv2.imshow("raw_img",res_d)
cv2.imshow("res",res_a)
cv2.imwrite("./save/res.png",res_a)
cv2.waitKey(0)
